package sql

import (
	"database/sql"
	"database/sql/driver"
	"fmt"
	"io"
	"sync"
	"sync/atomic"

	_ "github.com/go-sql-driver/mysql"
	"golang.org/x/net/context"

	"xp/lib/errors"
	"xp/lib/etc"
	"xp/lib/finalize"
	"xp/lib/log"
	"xp/lib/trace"
)

var (
	ErrNoRows              = sql.ErrNoRows
	ErrInvalidRowsAffected = errors.New("invalid rows affected")
)

type DB struct {
	*sql.DB
	*stmt
}

func (db *DB) Close() error {
	Close(db.stmt)
	Close(db.DB)
	return nil
}

func Default(dsn string) (*DB, error) {
	return Open(etc.Str("sql", "driver"), dsn)
}

var (
	dsnLock = new(sync.Mutex)
	dsnMap  = make(map[string]*DB)
)

func Open(driver, dsn string) (*DB, error) {
	dsnLock.Lock()
	defer dsnLock.Unlock()
	if db, ok := dsnMap[dsn]; ok {
		return db, nil
	}

	db, err := sql.Open(driver, dsn)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(int(etc.Int("sql", "max_open_conns")))
	db.SetMaxIdleConns(int(etc.Int("sql", "max_idel_conns")))
	ndb := &DB{DB: db, stmt: newStmt()}
	dsnMap[dsn] = ndb
	finalize.OnExit(ndb.Close)
	return ndb, nil
}

func Close(closer io.Closer) {
	if err := closer.Close(); err != nil {
		log.Warn("err", err.Error())
	}
}

func Rollback(tx *Tx) {
	if err := tx.Rollback(); err != nil {
		log.Warn("err", err.Error())
	}
}

type StmtStore map[string]*sql.Stmt

type stmt struct {
	sync.Mutex
	atomic.Value
}

func newStmt() *stmt {
	s := new(stmt)
	s.Store(make(StmtStore))
	return s
}

func (s *stmt) GetStmt(db *sql.DB, sql string, force bool) (*sql.Stmt, error) {
	store := s.Load().(StmtStore)
	v, ok := store[sql]
	if ok && !force {
		return v, nil
	}
	if ok {
		v.Close()
	}

	stmt, err := db.Prepare(sql)
	if err != nil {
		return nil, err
	}
	s.Lock()
	store = s.Load().(StmtStore)
	n := make(StmtStore)
	for k, v := range store {
		n[k] = v
	}
	n[sql] = stmt
	s.Store(n)
	s.Unlock()
	return stmt, nil
}

func (s *stmt) Close() error {
	s.Lock()
	store := s.Load().(StmtStore)
	for k, _ := range store {
		Close(store[k])
	}
	s.Unlock()
	return nil
}

func (db *DB) Exec(ctx context.Context, sql string, affect int64, args ...interface{}) error {
	span := trace.NewSpan(ctx, "Exec", "sql", sql, "args", fmt.Sprintln(args...))
	err := db.exec(sql, affect, args...)
	if span != nil {
		span.Finish(err)
	}
	return errors.Wrap(err)
}

func (db *DB) exec(sql string, affect int64, args ...interface{}) error {
	var force bool
	for {
		stmt, err := db.GetStmt(db.DB, sql, force)
		if err != nil {
			return err
		}
		rs, err := stmt.Exec(args...)
		if err != nil {
			if force {
				return err
			}
			if err == driver.ErrBadConn {
				force = true
				continue
			}
			return err
		}
		if affect > 0 {
			affected, err := rs.RowsAffected()
			if err != nil {
				return err
			}
			if affected != affect {
				return ErrInvalidRowsAffected
			}
		}
		return nil
	}
}

func (db *DB) ExecX(ctx context.Context, sql string, affect int64, args ...interface{}) (int64, error) {
	span := trace.NewSpan(ctx, "ExecX", "sql", sql, "args", fmt.Sprintln(args...))
	id, err := db.execX(sql, affect, args...)
	if span != nil {
		span.Finish(err)
	}
	return id, errors.Wrap(err)
}

func (db *DB) execX(sql string, affect int64, args ...interface{}) (int64, error) {
	var force bool
	for {
		stmt, err := db.GetStmt(db.DB, sql, force)
		if err != nil {
			return 0, err
		}
		rs, err := stmt.Exec(args...)
		if err != nil {
			if force {
				return 0, err
			}
			if err == driver.ErrBadConn {
				force = true
				continue
			}
			return 0, err
		}
		if affect > 0 {
			affected, err := rs.RowsAffected()
			if err != nil {
				return 0, err
			}
			if affected != affect {
				return 0, ErrInvalidRowsAffected
			}
		}
		return rs.LastInsertId()
	}
}

func (db *DB) QueryRow(ctx context.Context, sql string, args ...interface{}) *Row {
	span := trace.NewSpan(ctx, "QueryRow", "sql", sql, "args", fmt.Sprintln(args...))
	rows, err := db.query(sql, args...)
	if span != nil {
		span.Finish(err)
	}
	return &Row{rows: rows, err: errors.Wrap(err)}
}

func (db *DB) Query(ctx context.Context, sql string, args ...interface{}) (*sql.Rows, error) {
	span := trace.NewSpan(ctx, "Query", "sql", sql, "args", fmt.Sprintln(args...))
	rows, err := db.query(sql, args...)
	if span != nil {
		span.Finish(err)
	}
	return rows, errors.Wrap(err)
}

func (db *DB) query(sql string, args ...interface{}) (*sql.Rows, error) {
	var force bool
	for {
		stmt, err := db.GetStmt(db.DB, sql, force)
		if err != nil {
			// fallback
			return db.DB.Query(sql, args...)
		}
		rows, err := stmt.Query(args...)
		if err != nil {
			if force {
				return nil, err
			}
			if err == driver.ErrBadConn {
				force = true
				continue
			}
			return nil, err
		}
		return rows, nil
	}
}

func (db *DB) Begin(ctx context.Context) (*Tx, error) {
	span := trace.NewSpan(ctx, "tx")
	tx, err := db.DB.Begin()
	if err != nil {
		span.Finish(err)
		return nil, errors.Wrap(err)
	}
	t := txPool.Get().(*Tx)
	t.Tx = tx
	t.span = span
	return t, nil
}

var txPool = sync.Pool{
	New: func() interface{} {
		return new(Tx)
	},
}

type Tx struct {
	*sql.Tx
	span *trace.Span
}

func (tx *Tx) Commit() error {
	err := tx.Tx.Commit()
	if tx.span != nil {
		tx.span.Finish(err)
	}
	tx.Tx = nil
	tx.span = nil
	txPool.Put(tx)
	return errors.Wrap(err)
}

func (tx *Tx) Exec(sql string, affect int64, args ...interface{}) error {
	rs, err := tx.Tx.Exec(sql, args...)
	if err != nil {
		return errors.Wrap(err)
	}
	if affect > 0 {
		affected, err := rs.RowsAffected()
		if err != nil {
			return errors.Wrap(err)
		}
		if affected != affect {
			return errors.Wrap(ErrInvalidRowsAffected)
		}
	}
	return nil
}

func (tx *Tx) ExecX(sql string, affect int64, args ...interface{}) (int64, error) {
	rs, err := tx.Tx.Exec(sql, args...)
	if err != nil {
		return 0, errors.Wrap(err)
	}
	if affect > 0 {
		affected, err := rs.RowsAffected()
		if err != nil {
			return 0, errors.Wrap(err)
		}
		if affected != affect {
			return 0, errors.Wrap(ErrInvalidRowsAffected)
		}
	}
	return rs.LastInsertId()
}

func (tx *Tx) QueryRow(sql string, args ...interface{}) *sql.Row {
	return tx.Tx.QueryRow(sql, args...)
}

func (tx *Tx) Rollback() error {
	err := tx.Tx.Rollback()
	if tx.span != nil {
		tx.span.SetKeyvals("roolback", "true")
		tx.span.Finish(err)
	}
	tx.Tx = nil
	tx.span = nil
	txPool.Put(tx)
	return errors.Wrap(err)
}
